import os
import xml.etree.ElementTree as ET

def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars

def convert_voc_to_yolo(voc_folder, output_folder, name_dict=None):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for xml_file in os.listdir(voc_folder):
        if xml_file.endswith('.xml'):
            tree = ET.parse(os.path.join(voc_folder, xml_file))
            root = tree.getroot()
            size = root.find('size')  # 图像尺寸
            w = int(size.find('width').text)  # 图像宽
            h = int(size.find('height').text)

            yolo_file = os.path.splitext(xml_file)[0] + '.txt'
            yolo_path = os.path.join(output_folder, yolo_file)

            with open(yolo_path, 'w') as f:
                for obj in root.findall('object'):
                    cls = obj.find('name').text
                    if name_dict is not None:
                        cls = name_dict[cls]
                    bbox = obj.find('bndbox')
                    x_center = (float(bbox.find('xmin').text) + float(bbox.find('xmax').text)) / 2 / w
                    y_center = (float(bbox.find('ymin').text) + float(bbox.find('ymax').text)) / 2 / h
                    width = (float(bbox.find('xmax').text) - float(bbox.find('xmin').text)) / w
                    height = (float(bbox.find('ymax').text) - float(bbox.find('ymin').text)) / h

                    yolo_line = f"{cls} {x_center} {y_center} {width} {height}\n"
                    f.write(yolo_line)


voc_folder = 'D:\\神经网络数据集\\头盔检测\\数据集1\\Safety Helmet Detection_datasets\\annotations'
output_folder = 'D:\\神经网络数据集\\头盔检测\\数据集1\\Safety Helmet Detection_datasets\\annotations_yolo2'
name_dict= {"helmet": "0", "head":"1", "person":"2"}
convert_voc_to_yolo(voc_folder, output_folder, name_dict)
